/******************************************************************/
/*                                                                */
/*  Winpooch : Windows Watchdog                                   */
/*  Copyright (C) 2004-2005  Benoit Blanchon                      */
/*                                                                */
/*  This program is free software; you can redistribute it        */
/*  and/or modify it under the terms of the GNU General Public    */
/*  License as published by the Free Software Foundation; either  */
/*  version 2 of the License, or (at your option) any later       */
/*  version.                                                      */
/*                                                                */
/*  This program is distributed in the hope that it will be       */
/*  useful, but WITHOUT ANY WARRANTY; without even the implied    */
/*  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR       */
/*  PURPOSE.  See the GNU General Public License for more         */
/*  details.                                                      */
/*                                                                */
/*  You should have received a copy of the GNU General Public     */
/*  License along with this program; if not, write to the Free    */
/*  Software Foundation, Inc.,                                    */
/*  675 Mass Ave, Cambridge, MA 02139, USA.                       */
/*                                                                */
/******************************************************************/


/******************************************************************/
/* Includes                                                       */
/******************************************************************/

// module's interface
#include "Disasm.h"

// project's headers
#include "Assert.h"
#include "Trace.h"


/******************************************************************/
/* Internal macros                                                */
/******************************************************************/

#define MOD(X)		(((X)>>6)&3)
#define OPCODE(X)	(((X)>>3)&7)
#define RM(X)		((X)&7)



/******************************************************************/
/* Internal function                                              */
/******************************************************************/

int _Dasm_ReadModRmAndSid (int nCur, BYTE * pInst)
{
  int nSize = 0 ;
  
  switch( MOD(pInst[nCur]) )
    {
    case 0: // Mod==00 => No disp

      // /!\ not valid if R/M is 4
      switch( RM(pInst[nCur]) )
	{
	case 5:
	  nSize = nCur + 9 ;
	  break ;
	default:
	  nSize = nCur + 1 ;
	}

      break ; 
      
    case 1: // Mod==01 => 8 bits disp

      // /!\ maybe not valid if R/M is 0
      switch( RM(pInst[nCur]) )
	{ 
	case 4: // an SIB follows
	  nSize = nCur + 3 ;
	  break ;

	default:
	  nSize = nCur + 2 ;
	}

      break ;
      
    case 2: // Mod==10 => 32 bits disp

      // /!\ not valid if R/M is 4
      nSize = nCur + 5 ;
      
      break ;
      
    case 3:
      nSize = nCur + 1 ;
      break ; 
    }

  return nSize ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

int Dasm_GetInstructionSize (BYTE * pInst) 
{
  int nSize = 0 ;

  switch( pInst[0] )
    {
    case 0x51:	// PUSH ECX
    case 0x55:	// PUSH EBP
      nSize = 1 ;
      break ;

    case 0x68:	// PUSH imm32
    case 0xB8:  // MOV eax, imm32
      nSize = 5 ;
      break ;  

    case 0x6A:	// PUSH imm8
      nSize = 2 ;
      break ;

    case 0x83:      
      switch( OPCODE(pInst[1]) )
	{
	case 5: // SUB r/m32,imm8
	  nSize = 3 ;
	  break ;
	}
      break ;

    case 0x33:	// XOR r32,r/m32
    case 0x39:	// CMP r/m32,r32
    case 0x81:  // /7  CMP r/m32,r/m32
		// /5  SUB r/m32,imm32
    case 0x85:	// TEST r/m32,r32
    case 0x8B:	// MOV r32,r/m32
      nSize = _Dasm_ReadModRmAndSid (1, pInst) ;
      break ;

    case 0x8D:	// LEA	r32,m
      nSize = 4 ;
      break ;

    case 0xE8: // CALL rel32
    case 0xE9: // JMP rel32
      TRACE_WARNING (TEXT("Relative instruction found (0x%02X)\n"), pInst[0]) ;
      nSize = 5 ;
      break ;

    case 0xFF:	
      switch( OPCODE(pInst[1]) )
	{
	case 6: // PUSH r/m32
	  nSize = _Dasm_ReadModRmAndSid (1, pInst) ;
	  break ;
	}
      break ;

    default:
      TRACE_ERROR (TEXT("Unknown instruction 0x%02X\n"), pInst[0]) ;
      break ; 
    }

  return nSize ;
}


/******************************************************************/
/* Exported function                                              */
/******************************************************************/

BOOL  Dasm_MoveProgram (BYTE * pDst, int nDstMax, 
			BYTE * pSrc, int nSrcSize)
{
  int		nReadBytes = 0 ;
  int		nInstSize ;
  BYTE		*pInst ;
  DWORD		dwAddress ;

  if( nDstMax < nSrcSize ) {
    TRACE_ERROR (TEXT("Destination buffer too small\n")) ;
    return FALSE ;
  }

  while( nReadBytes < nSrcSize )
    {
      pInst = pSrc + nReadBytes ;
      nInstSize = Dasm_GetInstructionSize (pInst) ;
      
      if( ! nInstSize ) {
	TRACE_ERROR (TEXT("GetInstructionSize failed\n")) ;
	return FALSE ;
      }
      
      if( nReadBytes+nInstSize > nSrcSize ) {
	TRACE_ERROR (TEXT("Invalid source size\n")) ;
	return FALSE ;
      }

      switch( pInst[0] )
	{ 
	case 0xE8: // CALL rel32
	case 0xE9: // JMP rel32
	  
	  dwAddress = *(DWORD*)(pInst+1) ;		// get source relative
	  dwAddress += (DWORD)pSrc + nReadBytes ;	// convert to absolute
	  dwAddress -= (DWORD)pDst + nReadBytes ;	// convert to relative
	  pDst[nReadBytes] = pInst[0] ;
	  *((DWORD*)&pDst[nReadBytes+1]) = dwAddress ;

	  TRACE_INFO (TEXT("Relative instruction patched (0x%02X)\n"), pInst[0]) ;

	  break ;
	  
	default:
	  memcpy (pDst+nReadBytes, pSrc+nReadBytes, nInstSize) ;
	}    
      
      nReadBytes += nInstSize ;
    }
  
  ASSERT (nReadBytes==nSrcSize) ;

  return TRUE ;
}
